package syncx

import (
	"github.com/DiracLee/dires-go/utils"
	"sync"
)

type LockBucket interface {
	Lock(keys ...string)
	Unlock(keys ...string)
	RLock(keys ...string)
	RUnlock(keys ...string)
	RWLock(readKeys, writeKeys []string)
	RWUnlock(readKeys, writeKeys []string)
}

type lockBucket struct {
	locks []*sync.RWMutex
}

func NewLockBucket(size int) LockBucket {
	locks := make([]*sync.RWMutex, size)
	for i := 0; i < size; i++ {
		locks[i] = &sync.RWMutex{}
	}
	return &lockBucket{locks: locks}
}

func (b *lockBucket) Lock(keys ...string) {
	indices := utils.GetIndices(len(b.locks), keys)
	for _, index := range indices {
		b.locks[index].Lock()
	}
}

func (b *lockBucket) Unlock(keys ...string) {
	indices := utils.GetIndices(len(b.locks), keys)
	for _, index := range indices {
		b.locks[index].Unlock()
	}
}

func (b *lockBucket) RLock(keys ...string) {
	indices := utils.GetIndices(len(b.locks), keys)
	for _, index := range indices {
		b.locks[index].RLock()
	}
}

func (b *lockBucket) RUnlock(keys ...string) {
	indices := utils.GetIndices(len(b.locks), keys)
	for _, index := range indices {
		b.locks[index].RUnlock()
	}
}

func (b *lockBucket) RWLock(readKeys, writeKeys []string) {
	keys := append(readKeys, writeKeys...)
	indices := utils.GetIndices(len(b.locks), keys)
	writeIndexSet := utils.GetIndexSet(len(b.locks), writeKeys)
	for _, index := range indices {
		_, exists := writeIndexSet[index]
		mu := b.locks[index]
		if exists {
			mu.Lock()
			continue
		}
		mu.RLock()
	}
}

func (b *lockBucket) RWUnlock(readKeys, writeKeys []string) {
	keys := append(readKeys, writeKeys...)
	indices := utils.GetIndices(len(b.locks), keys)
	writeIndexSet := utils.GetIndexSet(len(b.locks), writeKeys)
	for _, index := range indices {
		_, exists := writeIndexSet[index]
		mu := b.locks[index]
		if exists {
			mu.Unlock()
			continue
		}
		mu.RUnlock()
	}
}
